#include "solver.h"
#include <iostream>
#include <fstream>

using std::cout;

Solver::Solver(int nx, int ny, int nz) 
  : _nx(nx), _ny(ny), _nz(nz), 
    _T(_nx*_ny*_nz), 
    _T_next(_nx*_ny*_nz),
    _stencil(_nx*_ny*_nz)
{    
    cout << _T.size() << " degrees of freedom.\n";

    // TODO - allow alpha to be set.
    _alpha = 1.0e-4;

    double dx = 1.0/double(_nx-1);
    double dy = 1.0/double(_ny-1);
    double dz = 1.0/double(_nz-1);
    _idx2 = 1.0/(dx*dx);
    _idy2 = 1.0/(dy*dy);
    _idz2 = 1.0/(dz*dz);

    // For each point, build the stencil.
    for (int i=0; i<size(); ++i) {
        IJK ijk = to_ijk(i);

        _stencil[i].xm = i - 1;
        _stencil[i].xp = i + 1;
        _stencil[i].ym = i - _nx;
        _stencil[i].yp = i + _nx;
        _stencil[i].zm = i - _nx*_ny;
        _stencil[i].zp = i + _nx*_ny;
        if (ijk.i == 0)       _stencil[i].xm = _stencil[i].xp;
        if (ijk.i == _nx-1)   _stencil[i].xp = _stencil[i].xm;
        if (ijk.j == 0)       _stencil[i].ym = _stencil[i].yp;
        if (ijk.j == _ny-1)   _stencil[i].yp = _stencil[i].ym;
        if (ijk.k == 0)       _stencil[i].zm = _stencil[i].zp;
        if (ijk.k == _nz-1)   _stencil[i].zp = _stencil[i].zm;
    }   
}

// Set initial conditions.
void Solver::init() {
    for (int i=0; i<size(); ++i) {
        IJK ijk = to_ijk(i);
        if (ijk.i == 0)  _T(i) = _T_next(i)= 1.0;
        else             _T(i) = _T_next(i) = 0.0;
    }
}

void Solver::fwd_euler_step(int nsteps, double dt) {
    init();
    std::fstream fid("soln_fwd_euler", std::ios::out);
    for (int step=0; step<nsteps; ++step) {
        for (int i=0; i<size(); ++i) {

            IJK ijk = to_ijk(i);
            if (ijk.i == 0 || ijk.i == _nx-1) continue;

            LaplacianStencil &s = _stencil[i];

            double d2x = (_T(s.xm) + _T(s.xp) - 2.0*_T(i)) * _idx2;
            double d2y = (_T(s.ym) + _T(s.yp) - 2.0*_T(i)) * _idy2;
            double d2z = (_T(s.zm) + _T(s.zp) - 2.0*_T(i)) * _idy2;

            _T_next(i) = _T(i) + _alpha*dt * (d2x + d2y + d2z);

        }
        _T_next.swap(_T);
        
        if (step % 10 > 0) continue;
	    for (int i=0; i<_nx; ++i) {
            IJK ijk = {i,0,0};
		    fid << _T(from_ijk(ijk));
            if (i==_nx-1) fid << "\n";
            else          fid << " ";
    	}
    }
}

void Solver::cn_step(int nsteps, double dt) {
    init();
    std::fstream fid("soln_cn", std::ios::out);

    const int max_iter = 50;
    const double c = 0.5*_alpha*dt;

    Vector b(size()), xk1(size()), xk0(size());
    for (int step=0; step<nsteps; ++step) {
        for (int i=0; i<size(); ++i) {
            IJK ijk = to_ijk(i);
            if (ijk.i == 0 || ijk.i == _nx-1) {
                xk0(i) = xk1(i) = _T(i);
                continue;
            }

            LaplacianStencil &s = _stencil[i];
            double d2x = (_T(s.xm) + _T(s.xp) - 2.0*_T(i)) * _idx2;
            double d2y = (_T(s.ym) + _T(s.yp) - 2.0*_T(i)) * _idy2;
            double d2z = (_T(s.zm) + _T(s.zp) - 2.0*_T(i)) * _idy2;
            double d2 = c*(d2x + d2y + d2z);

            b(i)  = _T(i) + d2;        // Kind of a half step.
            _T_next(i) = _T(i) + d2;   // Guess of full step.

        }
        double aii = 1.0 + 2.0*c*(_idx2+_idy2+_idz2);
        double inv_aii = 1.0 / aii;
        // Gauss-Seidel iteration on k.
        for (int k=0;; ++k) {
            double err=0.0, err0;
            for (int i=0; i<size(); ++i) {
                IJK ijk = to_ijk(i);
                if (ijk.i == 0 || ijk.i == _nx-1) continue;

                LaplacianStencil &s = _stencil[i];
                // axk is a_ij * x_j for j > i
                double axk0 = -c*(xk0(s.xp)*_idx2
                                 +xk0(s.yp)*_idy2
                                 +xk0(s.zp)*_idz2);

                // axk1 is a_ij * x_j for j < i
                double axk1 = -c*(xk1(s.xm)*_idx2
                                 +xk1(s.ym)*_idy2
                                 +xk1(s.zm)*_idz2);
                xk1(i) = inv_aii * (b(i) - axk0 - axk1);
                double e = (xk0(i)-xk1(i));
                err += e*e;
            }
            if (k==0) err0 = err;
            else if (err < 1e-8*err0) {
                break;
            }
            else if (k == max_iter-1) {
                cout << "Solution did not converge after " << k << " steps.\n";
                break;
            }
            xk1.swap(xk0);
        }

        // Transfer solution to _T
        for (int i=0; i<size(); ++i) _T(i) = xk1(i);

        if (step % 10 > 0) continue;
	    for (int i=0; i<_nx; ++i) {
            IJK ijk = {i,0,0};
		    fid << _T(from_ijk(ijk));
            if (i==_nx-1) fid << "\n";
            else          fid << " ";
    	}
    }
}

